import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import seaborn as sns
from sklearn import datasets
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets("MNIST_data/")

U, S, V = np.linalg.svd(mnist.test.images)
print("S:", S)
k = 128
P = V[:k, :]
z = np.dot(mnist.test.images, P.T)
print("z:",z.shape)
pca = z

Set = sns.color_palette("Set2", 10)  # 调色板，10种颜色
color_mapping = {key: value for (key, value) in enumerate(Set)}  # color_mapping创建颜色字典
colors = list(map(lambda x: color_mapping[x], mnist.test.labels))  # 将每个样本的标记映射到相应的颜色
fig = plt.figure()
ax = Axes3D(fig)
ax.scatter(pca[:, 0], pca[:, 1], pca[:, 2], c=colors)  # 绘制散点图,用前三个特征描述十个数字

plt.show()
